# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from collections import namedtuple
from experimental.prophet.graph import graph

_Level = namedtuple('Level', ['price', 'qty'])


class _Book:
  def __init__(self, fastfeed_obj, inverse_product=False):
    self._fastfeed_book = fastfeed_obj.book
    self._fastfeed_trade = fastfeed_obj.trade
    self._timestamp = None
    self._ask0 = None
    self._bid0 = None
    self._est_ask0_p = None
    self._est_bid0_p = None
    self._inverse = inverse_product

  def _ensure_best_level(self):
    if self._ask0 is None:
      assert self._bid0 is None
      with graph.control_if(None):
        outs = graph.create_node('fastfeed.BookBestLevel', [self._fastfeed_book]).outputs
      self._timestamp = outs.timestamp
      self._ask0 = _Level(outs.ask0_price, outs.ask0_qty)
      self._bid0 = _Level(outs.bid0_price, outs.bid0_qty)

  def _ensure_est_best_level(self):
    if self._est_ask0_p is None:
      assert self._est_bid0_p is None
      with graph.control_if(None):
        outs = graph.create_node('fastfeed.EstimatedBestLevel',
                                 [self._fastfeed_book, self._fastfeed_trade]).outputs
      self._est_ask0_p = outs.ask0_estimated_price
      self._est_bid0_p = outs.bid0_estimated_price

  @property
  def timestamp(self):
    self._ensure_best_level()
    return self._timestamp

  @property
  def ask0(self):
    self._ensure_best_level()
    return self._ask0

  @property
  def bid0(self):
    self._ensure_best_level()
    return self._bid0

  @property
  def estimated_ask0_price(self):
    self._ensure_est_best_level()
    return self._est_ask0_p

  @property
  def estimated_bid0_price(self):
    self._ensure_est_best_level()
    return self._est_bid0_p

  def query_bid_qty(self, price):
    res = graph.create_node('fastfeed.BookQtyQuery',
                            [self._fastfeed_book, graph.as_variable(price)],
                            attrs={
                                'query_bid': True, 'inverse_product': self._inverse
                            }).outputs
    return res

  def query_ask_qty(self, price):
    res = graph.create_node('fastfeed.BookQtyQuery',
                            [self._fastfeed_book, graph.as_variable(price)],
                            attrs={
                                'query_bid': False, 'inverse_product': self._inverse
                            }).outputs
    return res


def get_book_from_fastfeed(fastfeed_obj, **kwargs):
  cache = graph.get_aux_entry('book_from_fastfeed')
  key = '%s|%s' % (fastfeed_obj.book.instance_name, str(kwargs))
  if key not in cache:
    cache[key] = _Book(fastfeed_obj, **kwargs)
  return cache[key]
